package trace

import (
	"context"
	"fmt"
	"github.com/go-kratos/kratos/v2/metadata"
	"github.com/go-kratos/kratos/v2/transport"
	"github.com/google/uuid"
	"go-kratos-demo/internal/lib/contextx"
	"strings"
)

const (
	HeaderTraceId            = "trace-id"            //传递的traceId
	HeaderStressTestingKey   = "xes-request-type"    //压测标志的key
	HeaderStressTestingValue = "performance-testing" //压测标志真正的value
	XesTracePtsPrefix        = "pts_"                //压测traceID前缀
)

func NewTraceId() string {
	return "admin-" + uuid.NewString()
}

// IsStressTest 判断是否是压测
func IsStressTest(ctx context.Context) bool {
	//1、指定的 header  k=>v
	stressTestString := ctx.Value(HeaderStressTestingKey)
	if stressTestString != nil && fmt.Sprintf("%v", stressTestString) == HeaderStressTestingValue {
		return true
	}
	//2、traceId 包含压测前缀 XES_TRACE_PTS_PREFIX
	if md, ok := metadata.FromServerContext(ctx); ok {
		traceId := md.Get(HeaderTraceId)
		if traceId != "" && strings.HasPrefix(traceId, XesTracePtsPrefix) {
			return true
		}
	}

	return false
}

// GetTraceIdFromHeader 获取traceId
func GetTraceIdFromHeader(transport transport.Transporter) (traceId string) {
	traceId = transport.RequestHeader().Get(HeaderTraceId)
	if traceId == "" {
		traceId = transport.RequestHeader().Get(contextx.HeaderTraceId)
	}
	if traceId == "" {
		traceId = NewTraceId()
	}
	return
}
